`timescale 1ns / 1ps
//////////////////////////////////////////////////////////////////////////////////
// Company: 
// Engineer: 
// 
// Create Date: 2025/10/02 11:56:56
// Design Name: 
// Module Name: fc_2
// Project Name: 
// Target Devices: 
// Tool Versions: 
// Description: 
// 
// Dependencies: 
// 
// Revision:
// Revision 0.01 - File Created
// Additional Comments:
// 
//////////////////////////////////////////////////////////////////////////////////


module fc_2(
    input  wire         clk_i,
    input  wire         rst_n_i,
    input  wire         start_i,
    output wire         done_o,

    output reg [1:0]     input_data_addr_o,    
    input  wire [1023:0]   input_data_i,

    output reg  [6:0]    weight_addr_o,         // {out_channel[4:0], group_idx[1:0]}
    input  wire [255:0]  weight_i,              // 32 * 8-bit

    output reg  [4:0]    bias_addr_o,
    input  wire [7:0]    bias_i,                // 8-bit signed

    output reg           fc_output_wren_o,
    output reg  [47:0]   fc_output_data_o,
    output reg  [4:0]    fc_output_addr_o,  

    output wire [1023:0] mul_data1_o, // 输入数据, 32*32
    output wire [255:0]  mul_data2_o, // 权重数据, 32*8
    input  wire [1279:0] mul_result_i // 乘法结果, 32 * 42-bit (32x8 -> 40, 扩展到42)
    );

    parameter INPUT_DIM    = 128;
    parameter OUTPUT_DIM   = 32;
    parameter GROUP_SIZE   = 32;
    parameter GROUPS       = 4;
    parameter MULT_LATENCY = 3;
    parameter RAM_LATENCY  = 3;

    reg [4:0] out_channel;
    reg [1:0] group_idx;
    reg       running;

    reg [4:0] out_channel_dly [0:24];
    reg [1:0] group_idx_dly   [0:24];
    reg [7:0] bias_reg  [0:24];

    wire [47:0] final_sum;

    localparam INDEX_ADDR      = 0;
    localparam INDEX_RAM_OUT   = INDEX_ADDR + RAM_LATENCY;
    localparam INDEX_MUL_OUT   = INDEX_RAM_OUT + MULT_LATENCY;
    localparam ADDER_TREE_STAGES = 3;
    localparam INDEX_ADDER_OUT = INDEX_MUL_OUT + ADDER_TREE_STAGES;
    localparam INDEX_ACCUM     = INDEX_ADDER_OUT + 1; // ! 这里为什么
    localparam INDEX_BIAS      = INDEX_ACCUM + 1;
    localparam INDEX_RELU      = INDEX_BIAS + 1;
    localparam INDEX_WRITE     = INDEX_RELU + 1;
    localparam INDEX_DONE      = INDEX_WRITE + 1;

    localparam BIAS_PIPE_OFFSET = (INDEX_BIAS > RAM_LATENCY) ? (INDEX_BIAS - RAM_LATENCY) : 0;

    reg  [47:0] channel_accum [0:OUTPUT_DIM-1];
    reg [47:0] accumulator;
    reg [47:0] relu_result;

    reg [INDEX_DONE:0] done_shift;

    integer i;

    always @(posedge clk_i or negedge rst_n_i) begin
        if (!rst_n_i) begin
            running <= 0;
            out_channel <= 0;
            group_idx <= 0;
        end else if (start_i) begin
            running <= 1;
        end else if (done_shift[INDEX_DONE]) begin
            running <= 0;
        end
    end

    always @(posedge clk_i or negedge rst_n_i) begin
        if (!rst_n_i) begin
            group_idx <= 0;
        end else if (running && group_idx == GROUPS - 1) begin
            group_idx <= 0;
        end else if (running) begin
            group_idx <= group_idx + 1;
        end
    end

    always @(posedge clk_i or negedge rst_n_i) begin
        if (!rst_n_i) begin
            out_channel <= 0;
        end else if (running && group_idx == GROUPS - 1) begin
            out_channel <= out_channel + 1;
        end
    end

    // 延迟链 (简单移位)
    integer d;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            for(d=0; d<=24; d=d+1) begin
                out_channel_dly[d] <= 0;
                group_idx_dly[d]   <= 0;
            end
        end else begin
            out_channel_dly[0] <= out_channel;
            group_idx_dly[0]   <= group_idx;
            for(d=1; d<=24; d=d+1) begin
                out_channel_dly[d] <= out_channel_dly[d-1];
                group_idx_dly[d]   <= group_idx_dly[d-1];
            end
        end
    end

    // 阶段1：地址生成
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            input_data_addr_o <= 0;
            weight_addr_o     <= 0;
            bias_addr_o       <= 0;
        end else if(running) begin
            input_data_addr_o <= group_idx;
            weight_addr_o     <= {out_channel, group_idx};
            bias_addr_o       <= out_channel;
        end
    end

    // 阶段2-4：RAM 输出寄存对齐 (3拍)
    // 偏置延迟链：bias_i 已包含 BRAM 3 拍延迟，这里再移位用于对齐到 INDEX_BIAS
    integer b;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            for(b=0;b<=24;b=b+1) bias_reg[b] <= 0;
        end else begin
            bias_reg[0] <= bias_i;
            for(b=1;b<=24;b=b+1)
                bias_reg[b] <= bias_reg[b-1];
        end
    end

    assign mul_data1_o = input_data_i;
    assign mul_data2_o = weight_i;

    adder_tree_2 u_adder_tree(
        .clk(clk_i),
        .rst_n(rst_n_i),
        .data(mul_result_i),
        .sum_out(final_sum)
    );

    // 阶段11：分组累加：在 INDEX_ACCUM 周期 final_sum 对应 group_idx_dly[INDEX_ACCUM]
    integer ca;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            for(ca=0; ca<OUTPUT_DIM; ca=ca+1)
                channel_accum[ca] <= 48'd0;
        end else begin
            // FIXME 数学上好像可以不用区分第一组？
            if(group_idx_dly[INDEX_ACCUM] == 0) begin
                // 第一组：写入
                channel_accum[out_channel_dly[INDEX_ACCUM]] <= final_sum;
            end else begin
                // 后续组：累加
                channel_accum[out_channel_dly[INDEX_ACCUM]] <= channel_accum[out_channel_dly[INDEX_ACCUM]] + final_sum;
            end
        end
    end

    // 阶段12：加偏置 (INDEX_BIAS) 仅在最后一组
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            accumulator <= 0;
        end else begin
            if(group_idx_dly[INDEX_BIAS] == (GROUPS-1)) begin
                accumulator <= channel_accum[out_channel_dly[INDEX_BIAS]] +
                               {{40{bias_reg[BIAS_PIPE_OFFSET][7]}}, bias_reg[BIAS_PIPE_OFFSET]};
            end else begin
                accumulator <= 0;
            end
        end
    end

    always @(posedge clk_i or negedge rst_n_i) begin
        if (!rst_n_i) begin
            relu_result <= 0;
        end else if (accumulator[47]) begin
            relu_result <= 0;
        end else begin
            relu_result <= accumulator;
        end
    end

    always @(posedge clk_i or negedge rst_n_i) begin
        if (!rst_n_i) begin
            fc_output_wren_o <= 0;
            fc_output_data_o <= 0;
            fc_output_addr_o <= 0;
        end else if (running) begin
            if(group_idx_dly[INDEX_WRITE] == (GROUPS-1)) begin
                fc_output_wren_o <= 1'b1;
                fc_output_addr_o <= out_channel_dly[INDEX_WRITE];
                fc_output_data_o <= relu_result;
            end else begin
                fc_output_wren_o <= 1'b0;
            end
        end
    end

    // 阶段15：done 产生：last_launch -> 在 INDEX_DONE 周期输出
    wire last_launch = running &&
                       (out_channel == OUTPUT_DIM-1) &&
                       (group_idx == GROUPS-1);
    integer ds;
    always @(posedge clk_i or negedge rst_n_i) begin
        if(!rst_n_i) begin
            done_shift <= 0;
        end else begin
            done_shift <= {done_shift[INDEX_DONE-1:0], last_launch};
        end
    end

    assign done_o = done_shift[INDEX_DONE];

endmodule